/*
* create by Jinghui Wang on 2025-2-13
 * This file heavily references the repository of Xiang Gao
 * https://github.com/gaoxiang12/slam_in_autonomous_driving/blob/master/src/common/math_utils.h
 */

#include <iostream>
#include <chrono>
#include <random>
#include <numeric>
#include <Eigen/Core>
#include <Eigen/SVD>
#include <glog/logging.h>

/**
 * @brief 计算平面拟合，使用SVD分解
 * @tparam DataT 数据类型
 * @param data 数据输入，形式为三维向量数组
 * @param plane_coeffs 平面方程系数，形式为四维向量，前三维表示法向量，最后一维表示截距
 * @param eps 误差容错阈值
 */
template <typename DataT>
bool FitPlane(std::vector<Eigen::Matrix<DataT, 3, 1>>& data, Eigen::Matrix<DataT, 4, 1>& plane_coeffs,
              double eps = 1e-2) {
    // 检查数据是否足够
    if (data.size() < 3) {
        std::cerr << "Data size is not enough!" << std::endl;
        return false;
    }

    // 矩阵形式最小二乘
    /* A 的样式如下
     * | x11,   x12,    x13,    1 |
     * | x21,   x22,    x23,    1 |
     * | ...,   ...,    ...,    1 |
     * | xn1,   xn2,    xn3,    1 |
     */
    // Eigen::MatrixX4d A(data.size()); // 用法有误
    Eigen::MatrixXd A(data.size(), 4);
    for (int i = 0; i < data.size(); ++i) {
        A.row(i).head(3) = data[i].transpose();
        A.row(i)[3] = static_cast<DataT>(1);
    }

    // 计算SVD
    Eigen::JacobiSVD<Eigen::MatrixXd> svd(A, Eigen::ComputeThinV); // 这里指定计算V矩阵
    plane_coeffs = svd.matrixV().rightCols(1);                      // V的最后一列就是最小特征值对应的特征向量

    // check error
    for (int i = 0; i < data.size(); ++i) {
        double error = plane_coeffs.head(3).dot(data[i]) + plane_coeffs[3];
        if (error * error > eps) {
            std::cerr << "Fitting error is too large!" << std::endl;
            return false;
        }
    }

    return true;
}


int main(int argc, char** argv) {
    // 初始化日志
    google::InitGoogleLogging(argv[0]);
    FLAGS_stderrthreshold = google::INFO;
    FLAGS_colorlogtostderr = true;
    LOG(INFO) << "testing plane fitting";

    // 创建数据
    Eigen::Vector4d true_plane_coeffs(1, 2, 3, 4);
    true_plane_coeffs.normalize(); // 平面方向不变，大小归一化

    // 随机点生成
    std::random_device rd;
    std::mt19937 gen(rd());
    std::uniform_real_distribution<> dis(0.0, 1.0);  // 均匀分布，用于生成 [0, 1] 之间的随机数
    std::normal_distribution<> noise_dis(0.0, 0.01); // 高斯噪声分布，标准差为 0.01

    const int num_tested_points = 100; // 测试点数
    std::vector<Eigen::Vector3d> points;

    for (int i = 0; i < num_tested_points; ++i) {
        Eigen::Vector3d p(dis(gen), dis(gen), dis(gen)); // 随机点

        // 生成随机点与原点的直线，交于平面，交点为p，加噪
        double n4 = -p.dot(true_plane_coeffs.head<3>()) / (true_plane_coeffs[3] + std::numeric_limits<double>::min());
        p = p / (n4 + std::numeric_limits<double>::min());
        p += Eigen::Vector3d(noise_dis(gen), noise_dis(gen), noise_dis(gen));

        points.emplace_back(p);

        // 验证点在平面上
        LOG(INFO) << "res of p: " << p.dot(true_plane_coeffs.head<3>()) + true_plane_coeffs[3];
    }

    Eigen::Vector4d estimated_plane_coeffs;

    auto start = std::chrono::steady_clock::now();
    auto result = FitPlane(points, estimated_plane_coeffs);
    auto end = std::chrono::steady_clock::now();
    auto duration = std::chrono::duration_cast<std::chrono::microseconds>(end - start);

    LOG(INFO) << "Time cost: " << duration.count() << " us"; // 我的机器上，25us左右

    if (result) {
        LOG(INFO) << "Estimated plane coefficients: " << estimated_plane_coeffs.transpose();
        LOG(INFO) << "True plane coefficients: " << true_plane_coeffs.transpose();
    } else {
        LOG(INFO) << "Plane fitting failed!";
    }

}